from scipy import poly1d, pi, absolute, angle, log10, real, imag, logspace
from scipy.signal import lti
from numpy import poly, array, concatenate, zeros
from tools import ishurwitz, freqint

class lcsblock(lti):
    """ A Linear Control System block. This class is an
        extension of the LTI systems class.
    """
    def __init__(self,*args,**kwords):
        """Initialize the LTI system using either:
           (numerator, denominator)
           (zeros, poles, gain)
           (A, B, C, D) -- state-space.
        """
        lti.__init__(self,*args,**kwords)
        self.states = self.A.shape[0]

    def isstable(self):
        return ishurwitz(poly(self.A))
    
    def freqresp(self, w=None):
        if w == None:
            w = freqint(self)
        s = 1.0j*w
        num = self.num.squeeze()  # Yes this is ugly. Somehow there are different depth levels possible for the numerator.
        return w, (poly1d(num)(s))/(poly1d(self.den)(s))

    def bode(self, w=None):
        w, freq_resp = self.freqresp(w)
        ampl = 20*log10(absolute(freq_resp))
        phase = 180*angle(freq_resp)/pi
        return w,ampl,phase

    def nyquist(self, w=None):
        if w == None:
            w = freqint(self)
            lowest = log10(min(w))
            highest = log10(max(w))
            w = concatenate((logspace(-50,lowest,50),w,logspace(highest,100,50)))            
        w, freq_resp = self.freqresp(w)
        re = real(freq_resp)
        im = imag(freq_resp)
        return w,re,im

    def __add__(self, other):
        if type(other) != type(self): raise ValueError, "Dont know how to add %s to lcsblock" % type(other)
        if (self.inputs != other.inputs) or (self.outputs != other.outputs):
            raise ValueError, "Systems do not have same amount of in/outputs."
        block1 = concatenate((self.A,zeros((self.states,other.states))), axis=1)
        block2 = concatenate((zeros((other.states,self.states)),other.A), axis=1)
        newA = concatenate((block1,block2), axis=0)
        newB = concatenate((self.B, other.B), axis=0)
        newC = concatenate((self.C, other.C), axis=1)
        newD = self.D + other.D
        return lcsblock(newA,newB,newC,newD)

if __name__ == '__main__':
    ##from IPython.Shell import IPShellEmbed

    ##sys = LCSblock(-2,1,1,1)
    sys = lcsblock(array([-150]),array([-10+10.0j, -10-10.0j]),10)
    print sys.isstable()
    w1,a,p = sys.bode()
    w2,r,i = sys.nyquist()
    
    from pylab import plot, semilogx
    semilogx(w1,a)
##    t,y = sys.step()
##    plot(t,y)
